import numpy as np
import matplotlib.pyplot as plt

from Solver import EpsilonGreedy, DecayingEpsilonGreedy, ThompsonSampling


class BernoulliBandit:
    """伯努利多臂老虎机,输入K表示拉杆个数"""

    def __init__(self, K):
        self.probs = np.random.uniform(
            size=K
        )  # 随机生成K个0～1的数,作为拉动每根拉杆的获奖概率

        self.best_idx = np.argmax(self.probs)  # 获奖概率最大的拉杆
        self.best_prob = self.probs[self.best_idx]  # 最大的获奖概率
        self.K = K

    def step(self, k):
        # 当玩家选择了k号拉杆后,根据拉动该老虎机的k号拉杆获得奖励的概率返回1（获奖）或0（未
        # 获奖）
        if np.random.rand() < self.probs[k]:
            return 1
        else:
            return 0


def plot_results(solvers, solver_names, with_actions=False):
    """生成累积懊悔随时间变化的图像。输入solvers是一个列表,列表中的每个元素是一种特定的策略。
    而solver_names也是一个列表,存储每个策略的名称"""
    for idx, solver in enumerate(solvers):
        time_list = range(len(solver.regrets))
        plt.plot(time_list, solver.regrets, label=solver_names[idx])
        # 绘制 solver.actions 随时间变化的散点图
        if with_actions:
            plt.scatter(
                time_list,
                solver.actions,
                label=f"{solver_names[idx]} Actions",
                marker=".",
                alpha=0.5,
                color="lightcoral",
            )
    plt.xlabel("Time steps")
    plt.ylabel("Cumulative regrets")
    plt.title("%d-armed bandit" % solvers[0].bandit.K)
    plt.legend()
    plt.show()


if __name__ == "__main__":
    np.random.seed(1)  # 设定随机种子,使实验具有可重复性
    K = 10
    bandit = BernoulliBandit(K)
    print("随机生成了一个%d臂伯努利老虎机" % K)
    print(
        "获奖概率最大的拉杆为%d号,其获奖概率为%.4f"
        % (bandit.best_idx, bandit.best_prob)
    )

    # np.random.seed(1)
    # epsilon_greedy_solver = EpsilonGreedy(bandit, epsilon=0.01)
    # epsilon_greedy_solver.run(5000)
    # print("epsilon-贪婪算法的累积懊悔为：", epsilon_greedy_solver.regret)
    # plot_results([epsilon_greedy_solver], ["EpsilonGreedy"])

    # np.random.seed(0)
    # epsilons = [1e-4, 0.01, 0.1, 0.25, 0.5]
    # epsilon_greedy_solver_list = [EpsilonGreedy(bandit, epsilon=e) for e in epsilons]
    # epsilon_greedy_solver_names = ["epsilon={}".format(e) for e in epsilons]
    # for solver in epsilon_greedy_solver_list:
    #     solver.run(5000)
    # plot_results(epsilon_greedy_solver_list, epsilon_greedy_solver_names)

    # np.random.seed(1)
    # decaying_epsilon_greedy_solver = DecayingEpsilonGreedy(bandit)
    # decaying_epsilon_greedy_solver.run(5000)
    # print(
    #     "epsilon值衰减的贪婪算法的累积懊悔为：", decaying_epsilon_greedy_solver.regret
    # )
    # plot_results(
    #     [decaying_epsilon_greedy_solver], ["DecayingEpsilonGreedy"], with_actions=True
    # )

    np.random.seed(1)
    thompson_sampling_solver = ThompsonSampling(bandit)
    thompson_sampling_solver.run(5000)
    print("汤普森采样算法的累积懊悔为：", thompson_sampling_solver.regret)
    plot_results([thompson_sampling_solver], ["ThompsonSampling"])
